from flask import Flask, render_template, request, jsonify, send_file, render_template, g
from flask_cors import CORS
from flask_login import LoginManager, login_required, current_user
from sqlalchemy.exc import SQLAlchemyError
from config import config
from database import db, Server, TransferTask, User
from utils.ssh_utils import ssh_utils
from datetime import datetime
import os
import threading
import docker
import logging
import time
from utils.docker_utils import DockerUtils
from auth import auth, init_admin_user
from sqlalchemy import text

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

# 初始化Flask-Login
login_manager = LoginManager()

# 全局Docker工具类实例，避免重复初始化
docker_utils_instance = None

def get_docker_utils():
    """获取Docker工具类实例（单例模式）"""
    global docker_utils_instance
    if docker_utils_instance is None:
        try:
            # 使用配置文件中的超时设置
            docker_utils_instance = DockerUtils()
            logger.info("Docker工具类初始化成功")
        except Exception as e:
            logger.error(f"Docker工具类初始化失败: {str(e)}")
            # 保留None状态，让调用者处理初始化失败的情况
    return docker_utils_instance

def create_app():
    app = Flask(__name__)
    app.config.from_object(config)
    
    # 如果是SQLite，增强引擎参数以减少锁冲突
    try:
        if str(app.config.get('SQLALCHEMY_DATABASE_URI', '')).startswith('sqlite:///'):
            app.config.setdefault('SQLALCHEMY_ENGINE_OPTIONS', {})
            engine_opts = app.config['SQLALCHEMY_ENGINE_OPTIONS']
            connect_args = engine_opts.get('connect_args', {})
            # 允许跨线程、增加超时
            connect_args.setdefault('check_same_thread', False)
            connect_args.setdefault('timeout', 30)
            engine_opts['connect_args'] = connect_args
            engine_opts.setdefault('pool_pre_ping', True)
    except Exception as e:
        logger.warning(f"设置SQLite引擎选项失败: {e}")
    
    # 初始化Flask-Login
    login_manager.init_app(app)
    login_manager.login_view = 'auth.login'
    login_manager.login_message = '请先登录才能访问此页面'
    login_manager.login_message_category = 'warning'
    
    # Initialize extensions
    db.init_app(app)
    CORS(app)
    
    # 注册认证蓝图
    app.register_blueprint(auth, url_prefix='/auth')
    
    # Create tables & PRAGMA优化
    with app.app_context():
        db.create_all()
        # SQLite PRAGMA 优化，开启WAL，延长busy_timeout，降低同步强度
        try:
            if str(app.config.get('SQLALCHEMY_DATABASE_URI', '')).startswith('sqlite:///'):
                db.session.execute(text('PRAGMA journal_mode=WAL;'))
                db.session.execute(text('PRAGMA busy_timeout=5000;'))  # 5秒
                db.session.execute(text('PRAGMA synchronous=NORMAL;'))
                db.session.commit()
                logger.info('SQLite PRAGMA 已应用: WAL, busy_timeout=5000, synchronous=NORMAL')
        except Exception as e:
            logger.warning(f"应用SQLite PRAGMA失败: {e}")
        
        # 初始化管理员用户
        init_admin_user()
    
    return app

app = create_app()

@login_manager.user_loader
def load_user(user_id):
    """加载用户"""
    return User.query.get(int(user_id))

@app.route('/')
@login_required
def index():
    return render_template('index.html')

@app.route('/transfer', methods=['GET', 'POST'])
@login_required
def transfer():
    if request.method == 'POST':
        data = request.get_json()
        image_name = data.get('image_name')
        target_server_id = data.get('target_server_id')
        target_path = data.get('target_path')
        
        # 验证输入
        if not all([image_name, target_server_id, target_path]):
            return jsonify({'error': 'Missing required parameters'}), 400
        
        # 获取目标服务器信息
        target_server = Server.query.get(target_server_id)
        if not target_server:
            return jsonify({'error': 'Target server not found'}), 404
        
        # 创建传输任务，关联当前用户
        task = TransferTask(
            image_name=image_name,
            target_server_id=target_server_id,
            target_path=target_path,
            status='pending',
            progress=0,
            message='任务已创建，等待执行',
            user_id=current_user.id  # 关联当前用户
        )
        db.session.add(task)
        db.session.commit()
        
        # 启动后台任务
        thread = threading.Thread(target=process_transfer_task, args=(task.id,))
        thread.daemon = True
        thread.start()
        
        return jsonify({'success': True, 'message': 'Transfer task started', 'task_id': task.id})
    
    servers = Server.query.all()
    return render_template('transfer.html', servers=servers)

@app.route('/transfer/status/<int:task_id>')
@login_required
def transfer_status(task_id):
    """获取传输任务状态"""
    task = TransferTask.query.get(task_id)
    if not task:
        return jsonify({'success': False, 'error': '任务不存在'}), 404
    
    # 检查任务是否属于当前用户
    if task.user_id != current_user.id and not current_user.is_admin:
        return jsonify({'success': False, 'error': '无权限访问此任务'}), 403
    
    return jsonify({
        'success': True,
        'task_id': task.id,
        'status': task.status,
        'progress': task.progress,
        'message': task.message,
        'created_at': task.created_at.isoformat() if task.created_at else None,
        'completed_at': task.completed_at.isoformat() if task.completed_at else None
    })

def process_transfer_task(task_id):
    with app.app_context():
        task = TransferTask.query.get(task_id)
        if not task:
            logger.error(f"任务 {task_id} 不存在")
            return
        
        logger.info(f"开始处理传输任务 {task_id}: {task.image_name}")
        
        try:
            # 更新任务状态为拉取中
            update_task_status(task, 'pulling', 10, '开始拉取Docker镜像...')
            logger.info(f"任务 {task_id}: 开始拉取镜像 {task.image_name}")

            # 获取Docker工具实例
            docker_utils = get_docker_utils()
            if not docker_utils:
                raise Exception("Docker工具类初始化失败")
            
            # 拉取镜像
            image, error = docker_utils.pull_image(task.image_name)
            if error:
                raise Exception(f'Failed to pull image: {error}')
            
            logger.info(f"任务 {task_id}: 镜像拉取成功")
            
            # 更新任务状态为保存中
            update_task_status(task, 'pulling', 30, '镜像拉取中...')            
            # 保存镜像为tar文件
            update_task_status(task, 'saving', 40, '准备保存镜像为tar文件...')

            temp_dir = app.config['UPLOAD_FOLDER']
            os.makedirs(temp_dir, exist_ok=True)
            
            # 生成安全的文件名
            safe_image_name = task.image_name.replace('/', '_').replace(':', '_')
            tar_filename = f"{safe_image_name}_{int(time.time())}.tar"
            tar_path = os.path.join(temp_dir, tar_filename)
            
            logger.info(f"任务 {task_id}: 开始保存镜像到 {tar_path}")
            saved_path, error = docker_utils.save_image(task.image_name, tar_path)
            if error:
                raise Exception(f'Failed to save image: {error}')
            
            logger.info(f"任务 {task_id}: 镜像保存成功")
            
            # 更新任务状态为传输中
            update_task_status(task, 'saving', 60, '镜像保存完成')

            # 传输到目标服务器
            update_task_status(task, 'transferring', 70, '准备传输到目标服务器...') 
            # 连接到目标服务器并传输文件
            target_server = task.target_server
            logger.info(f"任务 {task_id}: 连接到目标服务器 {target_server.host}")
            
            from config import config as app_config
            transfer_optimizer = docker_utils.transfer_optimizer
            
            # 优先使用优化器进行传输（支持rsync/HTTP/SCP）
            use_optimizer = transfer_optimizer is not None
            
            # 传输文件 - 进度回调
            def transfer_progress_callback(transferred, total):
                if total > 0:
                    progress = 70 + (transferred / total) * 25  # 70-95% 用于传输
                    update_task_status(task, 'transferring', progress, f'传输中: {transferred}/{total} bytes')

            success = False
            error = None
            
            if use_optimizer:
                logger.info(f"任务 {task_id}: 使用传输优化器发起传输（rsync/HTTP/SCP）")
                # rsync/HTTP/SCP 自动选择由配置控制
                opt_success, opt_result = transfer_optimizer.transfer_file_optimized(
                    tar_path,
                    target_server.host,
                    task.target_path,
                    target_server.username,
                    password=target_server.password,
                    key_path=target_server.ssh_key_path
                )
                
                if opt_success and isinstance(opt_result, dict) and opt_result.get('method') == 'http':
                    # HTTP模式：需要在远端执行下载命令
                    logger.info(f"任务 {task_id}: HTTP传输准备就绪，开始通过SSH在目标机器执行下载命令")
                    ssh_ok, ssh_err = ssh_utils.connect(
                        target_server.host,
                        target_server.port,
                        target_server.username,
                        target_server.password,
                        target_server.ssh_key_path
                    )
                    if not ssh_ok:
                        raise Exception(f'SSH connection failed for HTTP download: {ssh_err}')
                    download_cmd = opt_result.get('download_cmd')
                    output, err_output, exec_err = ssh_utils.execute_command(download_cmd, timeout=app_config.DOCKER_SAVE_TIMEOUT)
                    if exec_err:
                        success, error = False, f'HTTP下载命令执行失败: {exec_err} | {err_output or output}'
                    else:
                        success, error = True, None
                        logger.info(f"任务 {task_id}: 目标机HTTP下载完成")
                    ssh_utils.close()
                else:
                    success, error = opt_success, opt_result if not opt_success else (True, None)
            else:
                # 回退使用SCP（原行为）
                logger.info(f"任务 {task_id}: 传输优化器不可用，回退到SCP传输")
                ssh_ok, ssh_err = ssh_utils.connect(
                    target_server.host,
                    target_server.port,
                    target_server.username,
                    target_server.password,
                    target_server.ssh_key_path
                )
                if not ssh_ok:
                    raise Exception(f'SSH connection failed: {ssh_err}')
                success, error = ssh_utils.transfer_file(
                    tar_path,
                    task.target_path,
                    progress_callback=transfer_progress_callback
                )
                ssh_utils.close()

            if not success:
                raise Exception(f'File transfer failed: {error}')
            
            logger.info(f"任务 {task_id}: 文件传输成功")
            
            # 清理临时文件
            if os.path.exists(tar_path):
                os.remove(tar_path)
                logger.info(f"任务 {task_id}: 临时文件清理完成")
            
            # 更新任务状态为完成
            # 验证文件传输
            update_task_status(task, 'verifying', 96, '验证文件完整性...')
            
            # 可以添加文件验证逻辑，比如检查文件大小、MD5等
            
            # 更新任务状态为完成
            update_task_status(task, 'completed', 100, '传输完成成功')
            task.completed_at = datetime.utcnow()
            logger.info(f"任务 {task_id}: 传输任务完成")
 
        except Exception as e:
            error_msg = f'任务失败: {str(e)}'
            logger.error(f"任务 {task_id}: {error_msg}", exc_info=True)
            update_task_status(task, 'failed', task.progress, error_msg)
            task.completed_at = datetime.utcnow()
            
            # 清理临时文件（如果存在）
            if 'tar_path' in locals() and os.path.exists(tar_path):
                try:
                    os.remove(tar_path)
                    logger.info(f"任务 {task_id}: 失败后临时文件清理完成")
                except Exception as cleanup_error:
                    logger.error(f"任务 {task_id}: 清理临时文件失败: {cleanup_error}")
        
        finally:
            try:
                db.session.commit()
                logger.info(f"任务 {task_id}: 数据库更新完成")
            except Exception as db_error:
                logger.error(f"任务 {task_id}: 数据库更新失败: {db_error}")
                db.session.rollback()

def update_task_status(task, status, progress, message):
    """更新任务状态辅助函数"""
    task.status = status
    task.progress = progress
    task.message = message
    task.updated_at = datetime.utcnow()
    db.session.commit()

@app.route('/servers', methods=['GET', 'POST'])
@login_required
def manage_servers():
    if request.method == 'POST':
        try:
            # 获取表单数据
            name = request.form.get('name')
            host = request.form.get('host')
            port = int(request.form.get('port', 22))
            username = request.form.get('username')
            password = request.form.get('password')
            ssh_key_path = request.form.get('ssh_key_path')
            auth_type = request.form.get('auth_type', 'password')

            # 验证必填字段
            if not all([name, host, username]):
                return jsonify({'error': '缺少必填字段'}), 400

            # 创建服务器对象
            server = Server(
                name=name,
                host=host,
                port=port,
                username=username,
                password=password if auth_type == 'password' else None,
                ssh_key_path=ssh_key_path if auth_type == 'key' else None
            )

            db.session.add(server)
            db.session.commit()

            return jsonify({'message': '服务器添加成功'})

        except Exception as e:
            db.session.rollback()
            return jsonify({'error': str(e)}), 500

    # GET 请求处理
    servers = Server.query.all()
    return render_template('servers.html', servers=servers)

@app.route('/history')
@login_required
def transfer_history():
    # 只显示当前用户的任务，管理员可以看到所有任务
    if current_user.is_admin:
        tasks = TransferTask.query.order_by(TransferTask.created_at.desc()).all()
    else:
        tasks = TransferTask.query.filter_by(user_id=current_user.id).order_by(TransferTask.created_at.desc()).all()
    return render_template('history.html', tasks=tasks)

@app.route('/api/tasks/<int:task_id>')
@login_required
def get_task_status(task_id):
    task = TransferTask.query.get_or_404(task_id)
    # 检查权限
    if task.user_id != current_user.id and not current_user.is_admin:
        return jsonify({'error': '无权限访问此任务'}), 403
    return jsonify(task.to_dict())

@app.route('/api/servers')
@login_required
def get_servers():
    servers = Server.query.all()
    return jsonify([server.to_dict() for server in servers])

@app.route('/api/servers/<int:server_id>', methods=['DELETE'])
@login_required
def delete_server(server_id):
    try:
        server = Server.query.get_or_404(server_id)
        db.session.delete(server)
        db.session.commit()
        return jsonify({'message': '服务器删除成功'})
    except Exception as e:
        db.session.rollback()
        return jsonify({'error': str(e)}), 500

@app.route('/api/servers/<int:server_id>/test', methods=['GET'])
@login_required
def test_server_connection(server_id):
    try:
        server = Server.query.get_or_404(server_id)
        
        logger.info(f"测试服务器连接: {server.name} ({server.host}:{server.port})")
        
        # 测试SSH连接
        from utils.ssh_utils import SSHUtils
        ssh = SSHUtils()
        
        result = ssh.test_connection(
            host=server.host,
            port=server.port,
            username=server.username,
            password=server.password,
            key_path=server.ssh_key_path
        )
        
        logger.info(f"连接测试结果: {result}")
        
        return jsonify(result)
            
    except Exception as e:
        logger.error(f"测试连接异常: {e}")
        return jsonify({
            'success': False, 
            'message': f'测试异常: {str(e)}',
            'response_time': 0,
            'timestamp': datetime.utcnow().strftime('%Y-%m-%d %H:%M:%S')
        })

@app.route('/api/docker/status')
@login_required
def docker_status():
    try:
        # 获取Docker工具类实例
        docker_utils = get_docker_utils()
        if not docker_utils:
            return jsonify({
                'status': 'error',
                'message': '无法初始化Docker工具类，请检查Docker是否正确安装',
                'connection_details': {
                    'type': 'local',
                    'socket_path': '/var/run/docker.sock'
                }
            }), 500
        success, message = docker_utils.check_docker_connection()
         # 获取额外的版本信息（可选）
        version_success, version_info = docker_utils.get_docker_version()
        # 构建响应数据
        response = {
            'status': 'connected' if success else 'disconnected',
            'message': message,
            'connection_details': {
                'type': 'local',
                'socket_path': '/var/run/docker.sock',
                'docker_version': version_info if version_success else 'unknown'
            },
            'timestamp': docker_utils.last_check_time
        }
        
        # 根据连接状态返回适当的HTTP状态码
        http_status = 200 if success else 503
        return jsonify(response), http_status
    except Exception as e:
        logger.error(f"处理Docker状态请求时发生错误: {str(e)}", exc_info=True)
        return jsonify({
            'status': 'error',
            'message': f"服务器处理错误: {str(e)}",
            'connection_details': {
                'type': 'local',
                'socket_path': '/var/run/docker.sock'
            }
        }), 500

@app.route('/api/reload_config', methods=['POST'])
@login_required
def reload_config():
    """重新加载配置"""
    try:
        from config import config
        config.reload_config()
        logger.info("配置已重新加载")
        return jsonify({
            'success': True,
            'message': '配置已重新加载',
            'config': {
                'chunk_size': f"{config.TRANSFER_CHUNK_SIZE // (1024*1024)}MB",
                'compression': config.TRANSFER_USE_COMPRESSION,
                'compression_level': config.TRANSFER_COMPRESSION_LEVEL,
                'parallel': config.TRANSFER_USE_PARALLEL,
                'parallel_chunks': config.TRANSFER_PARALLEL_CHUNKS,
                'http_transfer': config.TRANSFER_USE_HTTP,
                'rsync': config.TRANSFER_USE_RSYNC
            }
        })
    except Exception as e:
        logger.error(f"重载配置失败: {str(e)}")
        return jsonify({
            'success': False,
            'message': f'重载配置失败: {str(e)}'
        }), 500

@app.route('/api/delete_task/<int:task_id>', methods=['DELETE'])
@login_required
def delete_task(task_id):
    """删除传输任务"""
    try:
        task = TransferTask.query.get(task_id)
        if not task:
            return jsonify({
                'success': False,
                'message': '任务不存在'
            }), 404
        
        # 检查权限：只有任务创建者或管理员可以删除
        if task.user_id != current_user.id and not current_user.is_admin:
            return jsonify({
                'success': False,
                'message': '没有权限删除此任务'
            }), 403
        
        # 如果任务正在运行，先停止
        if task.status in ['running', 'pending']:
            # 这里可以添加停止任务的逻辑
            logger.warning(f"尝试删除正在运行的任务: {task_id}")
        
        # 删除相关的临时文件
        try:
            if task.local_file_path and os.path.exists(task.local_file_path):
                os.unlink(task.local_file_path)
                logger.info(f"删除本地文件: {task.local_file_path}")
        except Exception as e:
            logger.warning(f"删除本地文件失败: {str(e)}")
        
        # 从数据库删除任务
        db.session.delete(task)
        db.session.commit()
        
        logger.info(f"用户 {current_user.username} 删除了任务 {task_id}")
        
        return jsonify({
            'success': True,
            'message': '任务已删除'
        })
        
    except Exception as e:
        logger.error(f"删除任务失败: {str(e)}")
        return jsonify({
            'success': False,
            'message': f'删除任务失败: {str(e)}'
        }), 500
        
if __name__ == '__main__':
    app.run(debug=True, host='0.0.0.0', port=5000)
